#include "ControlledSimulator.h"
#include <math/random.h>

ControlledRobotSimulator::ControlledRobotSimulator()
  :robot(NULL),oderobot(NULL),controller(NULL)
{
  controlTimeStep = 0.01;
}

void ControlledRobotSimulator::Init(Robot* _robot,ODERobot* _oderobot,RobotController* _controller)
{
  robot=_robot;
  oderobot=_oderobot;
  controller=_controller;
  oderobot->SetConfig(robot->q);
  oderobot->SetVelocities(robot->dq);
  command.actuators.resize(robot->drivers.size());
  if(controller) {
    controller->Reset();
  }
  curTime = 0;
  nextControlTime = 0;

}

void ControlledRobotSimulator::SimulateSensors()
{
  if(!controller) return;

  oderobot->GetConfig(sensors.q);
  oderobot->GetVelocities(sensors.dq);
  //TODO: contact sensors

  if(!sensors.hasJointPosition) { cout<<"Don't have joint position feedback"<<endl; sensors.q.setZero(); }
  if(!sensors.hasJointVelocity) { cout<<"Don't have joint velocity feedback"<<endl; sensors.dq.setZero(); }
  //cout<<"Variance: "<<sensors.qvariance<<endl;
  if(!sensors.qvariance.isEmpty()) {
    //cout<<"q: "<<sensors.qvariance<<endl;
    for(int i=0;i<sensors.q.n;i++)
      sensors.q(i) += RandGaussian()*Sqrt(sensors.qvariance(i));
  }
  if(!sensors.qresolution.isEmpty()) {
    //cout<<"q: "<<sensors.qresolution<<endl;
    for(int i=0;i<sensors.q.n;i++) {
      if(sensors.qresolution(i) > 0) {
	sensors.q(i) = round(sensors.q(i)/sensors.qresolution(i))*sensors.qresolution(i);
      }
    }
  }
  if(!sensors.dqvariance.isEmpty()) {
    //cout<<"dq: "<<sensors.qvariance<<endl;
    for(int i=0;i<sensors.q.n;i++)
      sensors.dq(i) += RandGaussian()*Sqrt(sensors.dqvariance(i));
  }
  if(!sensors.dqresolution.isEmpty()) {
    //cout<<"dq: "<<sensors.dqresolution<<endl;
    for(int i=0;i<sensors.q.n;i++) {
      if(sensors.dqresolution(i) > 0) {
	sensors.dq(i) = int(sensors.dq(i)/sensors.dqresolution(i)+0.5)*sensors.dqresolution(i);
      }
    }
  }
  //floating joints get no feedback
  for(size_t i=0;i<robot->joints.size();i++) {
    if(robot->joints[i].type == RobotJoint::Floating) {
      vector<int> indices;
      robot->GetJointIndices(i,indices);
      for(size_t j=0;j<indices.size();j++) {
	sensors.q(indices[j])=0;
	sensors.dq(indices[j])=0;
      }
    }
  }
  for(size_t i=0;i<sensors.accelerometers.size();i++) {
    RigidTransform T;
    Vector3 w,v;
    oderobot->GetBodyTransform(sensors.accelerometers[i].link,T);
    oderobot->GetBodyVelocity(sensors.accelerometers[i].link,w,v);
    sensors.accelerometers[i].accel.setZero();
    sensors.accelerometers[i].accumVel = cross(w,T.R*sensors.accelerometers[i].pos)+v;
    sensors.accelerometers[i].accumPos = T*sensors.accelerometers[i].pos;
  }
  for(size_t i=0;i<sensors.gyroSensors.size();i++) {
    RigidTransform T;
    Vector3 w,v;
    oderobot->GetBodyTransform(sensors.gyroSensors[i].link,T);
    oderobot->GetBodyVelocity(sensors.gyroSensors[i].link,w,v);
    sensors.gyroSensors[i].angvel = w;
    sensors.gyroSensors[i].rotation = T.R;
  }
  for(size_t i=0;i<sensors.gravitySensors.size();i++) {
    RigidTransform T;
    oderobot->GetBodyTransform(sensors.gravitySensors[i].link,T);
    sensors.gravitySensors[i].gravity = T.R*Vector3(0,0,-1);
  }
  //TODO: rotational noise?  accelerometer noise?
  //TODO: gravity noise
}

void ControlledRobotSimulator::GetCommandedConfig(Config& q) 
{
  Assert(command.actuators.size() == robot->drivers.size());
  for(size_t i=0;i<command.actuators.size();i++) {
    RobotJointDriver& d=robot->drivers[i];
    if(command.actuators[i].mode == ActuatorCommand::PID)
      robot->SetDriverValue(i,command.actuators[i].qdes);
    else
      FatalError("Can't get commanded config for non-config drivers");
  }
  q = robot->q;
}

void ControlledRobotSimulator::GetCommandedVelocity(Config& dq)
{ 
  Assert(command.actuators.size() == robot->drivers.size());
  for(size_t i=0;i<command.actuators.size();i++) {
    RobotJointDriver& d=robot->drivers[i];
    if(command.actuators[i].mode == ActuatorCommand::PID)
      robot->SetDriverValue(i,command.actuators[i].dqdes);
    else
      FatalError("Can't get commanded config for non-config drivers");
  }
  dq = robot->dq;
}

void ControlledRobotSimulator::GetSensedConfig(Config& q)
{
  oderobot->GetConfig(q);
}

void ControlledRobotSimulator::GetSensedVelocity(Config& dq)
{
  oderobot->GetVelocities(dq);
}

void ControlledRobotSimulator::Step(Real dt)
{
  Real endOfTimeStep = curTime + dt;
 
  if(controller) {
    //the controller update happens less often than the PID update loop
    if(nextControlTime < endOfTimeStep) {
      //simulate sensors
      SimulateSensors();
      //update controller
      controller->sensors = &sensors;
      controller->command = &command;
      controller->Update(controlTimeStep);
      nextControlTime += controlTimeStep;
    }

    //get torques
    Assert(command.actuators.size() == robot->drivers.size());
    for(size_t i=0;i<command.actuators.size();i++) {
      RobotJointDriver& d=robot->drivers[i];
      ActuatorCommand& cmd=command.actuators[i];
      switch(cmd.mode) {
      case ActuatorCommand::OFF:
	printf("Warning: actuator off?\n");
	break;
      case ActuatorCommand::TORQUE:
	printf("Warning: direct torque?\n");
	oderobot->AddDriverTorque(i,Clamp(cmd.torque,d.tmin,d.tmax));
	break;
      case ActuatorCommand::PID:
	{
	  //simulate low level errors in the PID loop?
	  Real q=oderobot->GetDriverValue(i);
	  Real dq=oderobot->GetDriverVelocity(i);
	  //do we want to use the sensor? probably not
	  //Real q=sensors.q(d.linkIndices[0]);
	  //Real dq=sensors.dq(d.linkIndices[0]);
	  /*
	  if(q < robot->qMin(link) || q > robot->qMax(link)) {
	    printf("Warning: joint angle %s out of bounds\n",robot->linkNames[link].c_str());
	    printf("q=%g, qmin=%g, qmax=%g (deg)\n",RtoD(q),RtoD(robot->qMin(link)),RtoD(robot->qMax(link)));
	    //getchar();
	  }
	  */
	  if(d.type == RobotJointDriver::Normal || d.type == RobotJointDriver::Translation || d.type == RobotJointDriver::Rotation) {
	    Real td=Clamp(cmd.GetPIDTorque(q,dq),d.tmin,d.tmax);
	    oderobot->AddDriverTorque(i,td);
	    //printf("%d: Current %g,%g, desired %g,%g, torque desired %g, clamped %g\n",i,q,dq,cmd.qdes,cmd.dqdes,cmd.GetPIDTorque(q,dq),td);
	  }
	  else if(d.type == RobotJointDriver::Affine) {
	    Real td=Clamp(cmd.GetPIDTorque(q,dq),d.tmin,d.tmax);
	    //printf("Current %g,%g, desired %g,%g, torque desired %g, clamped %g\n",q,dq,cmd.qdes,cmd.dqdes,cmd.GetPIDTorque(q,dq),td);
	    //getchar();
	    Vector tjoints(d.linkIndices.size());
	    Vector driverBasis(d.linkIndices.size());
	    robot->SetDriverValue(i,q);
	    robot->SetDriverVelocity(i,dq);
	    Real mechStiffness = 20;
	    Real mechDamping = 0.2;
	    Real mechMaxTorque = 2;
	    for(size_t j=0;j<d.linkIndices.size();j++) {
	      int link = d.linkIndices[j];
	      driverBasis[j] = d.affScaling[j]; //todo: should be a transmission parameter in the joint driver
	      tjoints[j] = mechStiffness*(robot->q(link)-oderobot->GetLinkAngle(link)) + mechDamping*(robot->dq(link)-oderobot->GetLinkVelocity(link));
	    }
	    tjoints.madd(driverBasis,-tjoints.dot(driverBasis)/driverBasis.normSquared());
	    if(tjoints.norm() > mechMaxTorque)
	      tjoints *= mechMaxTorque/tjoints.norm();
	    //cout<<"Stabilizing torques: "<<tjoints<<endl;
	    tjoints.madd(driverBasis,td);
	    //cout<<"Torques: "<<tjoints<<endl;
	    for(size_t j=0;j<d.linkIndices.size();j++) 
	      oderobot->AddLinkTorque(d.linkIndices[j],tjoints[j]);
	  }
	  else {
	    FatalError("Unknown driver type");
	  }

	  //advance PID controller
	  if(cmd.kI*cmd.iterm > d.tmax) { cmd.iterm = d.tmax/cmd.kI; }
	  else if(cmd.kI*cmd.iterm < d.tmin) { cmd.iterm = d.tmin/cmd.kI; }
	  cmd.IntegratePID(q,dt);
	}
	break;
      case ActuatorCommand::LOCKED_VELOCITY:
	//TODO: clamp to braking velocitiy?
	oderobot->SetDriverFixedVelocity(i,cmd.desiredVelocity,cmd.torque);
	break;
      }
    }
  }
  
  curTime = endOfTimeStep;
}

void ControlledRobotSimulator::UpdateRobot()
{
  oderobot->GetConfig(robot->q);
  oderobot->GetVelocities(robot->dq);
  robot->UpdateFrames();
}

bool ControlledRobotSimulator::ReadState(File& f)
{
  if(!ReadFile(f,curTime)) return false;
  if(!ReadFile(f,nextControlTime)) return false;
  if(!ReadFile(f,command)) return false;
  if(controller) {
    if(!controller->ReadState(f)) return false;
  }
  return true;
}

bool ControlledRobotSimulator::WriteState(File& f) const
{
  if(!WriteFile(f,curTime)) return false;
  if(!WriteFile(f,nextControlTime)) return false;
  if(!WriteFile(f,command)) return false;
  if(controller) {
    if(!controller->WriteState(f)) return false;
  }
  return true;
}
